import sys
import pickle
from pathlib import Path

import pandas as pd

BASE = Path(__file__).resolve().parents[1]

# Pick the latest pred_best_*.pkl
candidates = sorted((BASE).glob('pred_best_*.pkl'), key=lambda p: p.stat().st_mtime, reverse=True)
if not candidates:
    print('No pred_best_*.pkl found under', BASE)
    sys.exit(1)
pred_path = candidates[0]
print('Using prediction file:', pred_path)
run_id = pred_path.stem.replace('pred_best_', '')

# Load prediction
with open(pred_path, 'rb') as f:
    pred = pickle.load(f)

# Ensure DataFrame with score column
if isinstance(pred, pd.Series):
    pred = pred.to_frame('score')
else:
    if 'score' not in pred.columns:
        # if there is only one column, use it as score
        if len(pred.columns) == 1:
            pred.columns = ['score']
        else:
            # fallback: create score from first column
            pred = pred.rename(columns={pred.columns[0]: 'score'})

# Reset index to columns
if pred.index.nlevels == 2:
    # try to set names
    try:
        pred.index.set_names(['datetime', 'instrument'], inplace=True)
    except Exception:
        pass
pred_df = pred.reset_index()
# Normalize column names
pred_df = pred_df.rename(columns={'datetime': 'date'})

# Save full prediction CSV
out_csv = BASE / 'pred_best_full.csv'
pred_df.to_csv(out_csv, index=False)
print('Saved full prediction CSV:', out_csv, 'rows:', len(pred_df))

# Try to also export with label joined
artifacts_dir = None
for exp_dir in (BASE / 'mlruns').iterdir():
    if not exp_dir.is_dir() or exp_dir.name == '.trash':
        continue
    candidate = exp_dir / run_id / 'artifacts'
    if candidate.exists():
        artifacts_dir = candidate
        break

if artifacts_dir is None:
    print('Artifacts dir not found for label; skip joined CSV.')
    sys.exit(0)
label_path = artifacts_dir / 'label.pkl'
if not label_path.exists():
    print('label.pkl not found; skip joined CSV.')
    sys.exit(0)

with open(label_path, 'rb') as f:
    label = pickle.load(f)
if isinstance(label, pd.Series):
    label = label.to_frame('label')
if 'label' not in label.columns:
    if len(label.columns) == 1:
        label.columns = ['label']
    else:
        label = label.rename(columns={label.columns[0]: 'label'})

label_df = label.reset_index()
label_df = label_df.rename(columns={'datetime': 'date'})

joined = pd.merge(pred_df, label_df, on=['date', 'instrument'], how='inner')
out_joined_csv = BASE / 'pred_best_full_with_label.csv'
joined.to_csv(out_joined_csv, index=False)
print('Saved joined CSV (with label):', out_joined_csv, 'rows:', len(joined))